#include "solver.h"

#ifdef OPENMP
#include <omp.h>
#endif

typedef struct
{
    void *L;
    void *U;
    BLOCK *B;
    TILE *T;
    // not used?
    int cost;
} TASK;

int counter = 0;

int task_num = 0;
TASK *scheduler;
// 用于记录非零TILE块所在坐标
int *mark_L;
int *mark_U;

// 二分递归求解指定TILE的dgetrf
// s=start, e=end
void tile_dgetrf2(MATRIX *A, const int bk, const int tk, int s, int e)
{
    BLOCK *B = A->block[bk * A->n + bk];
    TILE *T = B->tile[tk * B->n + tk];

    if (s >= e)
        return;

    if (e - s == 1)
    {
        // pivot(A, bk, tk, s);
        return;
    }
    else
    {
        int mid = (s + e) / 2;
        int m1 = mid - s;
        int m2 = e - mid;

        tile_dgetrf2(A, bk, tk, s, mid);

        // compute L of tile
        double *L = &T->a[mid * T->n + s];
        int ldl = T->n;
        cblas_dtrsm(CblasRowMajor, CblasRight, CblasUpper, CblasNoTrans, CblasNonUnit, m2, m1, 1.0, &T->a[s * T->n + s],
                    T->n, L, ldl);

        // schur complement
        cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m2, T->n - mid, m1, -1.0, L, ldl, &T->a[s * T->n + mid],
                    T->n, 1.0, &T->a[mid * T->n + mid], T->n);

        // TEST: Colmajor
        // cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, m2, T->n - mid, m1, -1.0, L, ldl, &T->a[s * T->n + mid],
        //             T->n, 1.0, &T->a[mid * T->n + mid], T->n);           

        task_num = 0;
        for (int tj = tk + 1; tj < B->n; tj++)
        {
            if (!B->tile[tk * B->n + tj])
                continue;
            scheduler[task_num].T = B->tile[tk * B->n + tj];
            ++task_num;
        }
        for (int bj = bk + 1; bj < A->n; bj++)
        {
            B = A->block[bk * A->n + bj];
            if (!B)
                continue;
            for (int tj = 0; tj < B->n; tj++)
            {
                if (!B->tile[tk * B->n + tj])
                    continue;
                scheduler[task_num].T = B->tile[tk * B->n + tj];
                ++task_num;
            }
        }

        for (int p = 0; p < task_num; ++p)
        {
            T = scheduler[p].T;
            cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m2, T->n, m1, -1.0, L, ldl, &T->a[s * T->n], T->n,
                        1.0, &T->a[mid * T->n], T->n);
        }

        tile_dgetrf2(A, bk, tk, mid, e);
    }
}

// 分解对角tile并计算其所在行的U TILE。伪代码：
// index = (bk, tk)
// dgetrf(A[index, index])
// A[index, index+1:] /= A[index, index]
void tile_dgetrf(MATRIX *A, int bk, int tk)
{
    BLOCK *B = A->block[bk * A->n + bk];
    if (!B)
    {
        int m = get_bs(A, bk);
        B = A->block[bk * A->n + bk] = construct_B(m, m);
    }

    TILE *T = B->tile[tk * B->n + tk];
    if (!T)
    {
        int m = get_ts(A, bk, tk);
        T = B->tile[tk * B->n + tk] = construct_T(m, m);
        B->nonz++;
    }

    tile_dgetrf2(A, bk, tk, 0, T->n);
}

// L /= T
// 其中T仅取上三角部分
inline void tile_dtrsmRUNN(const TILE *T, TILE *L)
{
    cblas_dtrsm(CblasRowMajor, CblasRight, CblasUpper, CblasNoTrans, CblasNonUnit, L->m, L->n, 1.0, T->a, T->n, L->a,
                L->n);
}

// T -= LU
inline void tile_dgemm(TILE *L, TILE *U, TILE **T)
{

    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, L->m, U->n, L->n, -1.0, L->a, L->n, U->a, U->n, 1.0, (*T)->a,
                U->n);

    counter++; // mutex TODO:
}

// 对角block的LU分解以及对应U BLOCK行的计算
void block_dgetrf(MATRIX *A, int bk)
{
    BLOCK *B = A->block[bk * A->n + bk];
    if (!B)
    {
        int m = get_bs(A, bk);
        B = A->block[bk * A->n + bk] = construct_B(m, m);
    }

    TILE *T1 = NULL;
    TILE *T2 = NULL;
    BLOCK *B2 = NULL;

    // 该循环内每一次循环分解一个DIAG TILE，一行U TILE与一列D TILE
    // 以及计算一个DIAG BLOCK与一行U BLOCK的schur补
    for (int k = 0; k < B->n; k++)
    {
        // 计算DIAG TILE与其所在行的U TILE
        tile_dgetrf(A, bk, k);
        // 计算DIAG BLOCK内的一列TILE
        task_num = 0;
        for (int i = k + 1; i < B->m; i++)
        {
            if (!B->tile[i * B->n + k])
                continue;
            scheduler[task_num].L = B->tile[i * B->n + k];
            ++task_num;
        }
        for (int p = 0; p < task_num; p++)
        {
            tile_dtrsmRUNN(B->tile[k * B->n + k], (TILE *)scheduler[p].L);
        }

        // 计算DIAG BLOCK的schur补
        int mark_L_size = 0;
        int mark_U_size = 0;
        for (int i = k + 1; i < B->m; ++i)
        {
            if (!B->tile[i * B->n + k])
                continue;
            mark_L[mark_L_size] = i;
            ++mark_L_size;
        }
        for (int j = k + 1; j < B->n; j++)
        {
            if (!B->tile[k * B->n + j])
                continue;
            mark_U[mark_U_size] = j;
            ++mark_U_size;
        }
        task_num = 0;
        for (int i = 0; i < mark_L_size; ++i)
        {
            int i1 = mark_L[i];
            for (int j = 0; j < mark_U_size; ++j)
            {
                int j1 = mark_U[j];
                scheduler[task_num].L = B->tile[i1 * B->n + k];
                scheduler[task_num].U = B->tile[k * B->n + j1];

                if (!B->tile[i1 * B->n + j1])
                {
                    B->tile[i1 * B->n + j1] = construct_T(B->tile[i1 * B->n + k]->m, B->tile[k * B->n + j1]->n);
                    B->nonz++;
                }
                scheduler[task_num].T = B->tile[i1 * B->n + j1];
                ++task_num;
            }
        }
        // 计算U BLOCK行的schur补
        for (int bj = bk + 1; bj < A->n; bj++)
        {
            B2 = A->block[bk * A->n + bj];
            if (!B2)
                continue;

            mark_U_size = 0;
            for (int j = 0; j < B2->n; ++j)
            {
                if (!B2->tile[k * B2->n + j])
                    continue;
                mark_U[mark_U_size] = j;
                ++mark_U_size;
            }
            for (int i = 0; i < mark_L_size; ++i)
            {
                int i1 = mark_L[i];
                for (int j = 0; j < mark_U_size; ++j)
                {
                    int j1 = mark_U[j];

                    scheduler[task_num].L = B->tile[i1 * B->n + k];
                    scheduler[task_num].U = B2->tile[k * B2->n + j1];

                    if (!B2->tile[i1 * B2->n + j1])
                    {
                        B2->tile[i1 * B2->n + j1] = construct_T(B->tile[i1 * B->n + k]->m, B2->tile[k * B2->n + j1]->n);
                        B2->nonz++;
                    }
                    scheduler[task_num].T = B2->tile[i1 * B2->n + j1];
                    ++task_num;
                }
            }
        }

// openmp1
#pragma omp parallel for
        for (int p = 0; p < task_num; p++)
        {
            tile_dgemm((TILE *)scheduler[p].L, (TILE *)scheduler[p].U, &scheduler[p].T);
        }
    }
}

// L /= B
// 其中B仅取上三角部分
void block_dtrsmRUNN(const BLOCK *B, BLOCK *L)
{
    if (!B || !L)
        return;

    TILE *T1 = NULL;
    TILE *L1 = NULL;

    for (int j = 0; j < L->n; j++)
    {
        for (int k = 0; k < j; k++)
        {
            if (!B->tile[k * B->n + j])
                continue;
            for (int i = 0; i < L->m; i++)
            {
                if (!L->tile[i * L->n + k])
                    continue;

                L1 = L->tile[i * L->n + k];
                T1 = B->tile[k * B->n + j];
                if (!L->tile[i * L->n + j])
                {
                    L->tile[i * L->n + j] = construct_T(L->tile[i * L->n + k]->m, B->tile[k * B->n + j]->n);
                    L->nonz++;
                }
                tile_dgemm(L->tile[i * L->n + k], B->tile[k * B->n + j], &L->tile[i * L->n + j]);
            }
        }
        for (int i = 0; i < L->m; i++)
        {
            if (!L->tile[i * L->n + j])
                continue;
            tile_dtrsmRUNN(B->tile[j * B->n + j], L->tile[i * L->n + j]);
        }
    }
}

// B -= LU
void block_dgemm(const BLOCK *L, const BLOCK *U, BLOCK *B)
{
    if (!L || !U || !B)
        return;
    for (int k = 0; k < L->n; k++)
    {
        for (int i = 0; i < B->m; i++)
        {
            if (!L->tile[i * L->n + k])
                continue;
            for (int j = 0; j < B->n; j++)
            {
                if (!U->tile[k * U->n + j])
                    continue;
                if (!B->tile[i * B->n + j])
                {
                    B->tile[i * B->n + j] = construct_T(L->tile[i * L->n + k]->m, U->tile[k * U->n + j]->n);
                    B->nonz++;
                }
                tile_dgemm(L->tile[i * L->n + k], U->tile[k * U->n + j], &B->tile[i * B->n + j]);
            }
        }
    }
}

void dgetrf(MATRIX *A)
{
    if (!A)
        return;

#ifdef OPENMP
    int block_avg_bandwidth = A->bandwidth / (2 * A->block_size * A->tile_size) + 1;
    int num_threads = MIN(block_avg_bandwidth * block_avg_bandwidth / 2 + 1, omp_get_max_threads());

    printf("block avg bandwidth = %d, num_threads = %d\n", block_avg_bandwidth, num_threads);
    omp_set_num_threads(num_threads);
#endif

    int max_task = A->m * A->n * A->block_size * A->block_size;

    // task scheduler
    scheduler = (TASK *)my_malloc(max_task, sizeof(TASK));
    mark_L = (int *)my_malloc(A->block_size * 2, sizeof(int));
    mark_U = mark_L + A->block_size;

    task_num = 0;

    BLOCK *B = NULL;
    BLOCK *L = NULL;
    BLOCK *U = NULL;

    struct timeval timetag;

    for (int k = 0; k < A->n; k++)
    {
        /* compute DIAG and U */
        set_time(&timetag);
        block_dgetrf(A, k);
        A->time_U += get_time(&timetag);

        /* compute L */
        set_time(&timetag);
        task_num = 0;
        for (int i = k + 1; i < A->m; i++)
        {
            if (!A->block[i * A->n + k])
                continue;
            scheduler[task_num].B = A->block[k * A->n + k];
            scheduler[task_num].L = A->block[i * A->n + k];
            scheduler[task_num].cost = A->block[k * A->n + k]->nonz * A->block[i * A->n + k]->nonz / A->block_size;
            ++task_num;
        }

// openmp2
#pragma omp parallel
        {
#pragma omp single nowait
            {
                for (int p = 0; p < task_num; p++)
                {
#pragma omp task shared(p)
                    block_dtrsmRUNN((BLOCK *)scheduler[p].B, (BLOCK *)scheduler[p].L);
                }
            }
        }

        A->time_L += get_time(&timetag);

        /* compute Schur complement */
        set_time(&timetag);
        task_num = 0;
        for (int i = k + 1; i < A->m; i++)
        {
            if (!A->block[i * A->n + k])
                continue;
            for (int j = k + 1; j < A->n; j++)
            {
                if (!A->block[k * A->n + j])
                    continue;
                if (!A->block[i * A->n + j])
                {
                    A->block[i * A->n + j] = construct_B(A->block[i * A->n + k]->m, A->block[k * A->n + j]->n);
                }
                scheduler[task_num].L = A->block[i * A->n + k];
                scheduler[task_num].U = A->block[k * A->n + j];
                scheduler[task_num].B = A->block[i * A->n + j];
                scheduler[task_num].cost = A->block[i * A->n + k]->nonz * A->block[k * A->n + j]->nonz / A->block_size;
                ++task_num;
            }
        }

// openmp3
#pragma omp parallel
        {
#pragma omp single nowait
            {
                for (int p = 0; p < task_num; p++)
                {
#pragma omp task shared(p)
                    block_dgemm((BLOCK *)scheduler[p].L, (BLOCK *)scheduler[p].U, (BLOCK *)scheduler[p].B);
                }
            }
        }
        A->time_C += get_time(&timetag);

#if 0
		progress_bar("dgetrf", A->n, k);
#endif
    }

    A->num_tgemm = counter;

    // information(A);

    my_free(scheduler, max_task, sizeof(TASK));
    my_free(mark_L, A->block_size * 2, sizeof(int));
}
